package com.ruoyi.ai.factory.impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.ruoyi.ai.domain.vo.WxChatVo;
import com.ruoyi.ai.domain.vo.WxModelConfigVo;
import com.ruoyi.ai.enums.MfgEnum;
import com.ruoyi.ai.factory.ModelChatTemplate;
import com.ruoyi.ai.ws.ChatWebSocketManager;
import com.ruoyi.common.exception.ServiceException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Semaphore;

@Service
@Slf4j
public class QWenModelChatService extends ModelChatTemplate {
    private Generation gen;

    @Resource
    private ChatWebSocketManager chatWebSocketManager;

    @PostConstruct
    public void createGen() {
        gen = new Generation();
    }
    public QWenModelChatService() {
        super(MfgEnum.QWEN.getType());
    }

    @Override
    public String chatStream(List<WxChatVo> chatVoList, Long userId) {
        GenerationParam param = this.getBaseParam();
        List<Message> messageList = this.buildMessageList(chatVoList);
        param.setMessages(messageList);
        try {
            log.info("发送的请求为{}", param);
            // 同步信号量
            Semaphore semaphore = new Semaphore(0);
            // 结果拼接对象
            StringBuilder resultBuilder = new StringBuilder();
            // 流式调用
            gen.streamCall(param, new ResultCallback<GenerationResult>() {

                @Override
                public void onEvent(GenerationResult generationResult) {
                    String newMessage = generationResult.getOutput().getChoices().get(0).getMessage().getContent();
                    StringBuilder finalResBuilder = resultBuilder.append(newMessage);
                    //据说代码块不终止
                    // 通过websocket返回给前端
                    //chatWebSocketManager.sendOneMessage(userId, finalResBuilder.toString());
                    chatWebSocketManager.sendOneMessage(userId, newMessage);
                }

                // 结束或者报错需要释放同步信号量
                @Override
                public void onComplete() {
                    semaphore.release();
                }

                @Override
                public void onError(Exception e) {
                    semaphore.release();
                    log.error("通义千问运行出错, 报错栈如下");
                    Throwable t = e;
                    while (t != null) {
                        log.error(t.toString());
                        t = e.getCause();
                    }
                }
            });
            semaphore.acquire();
            //将最终结果返回
            //发送对话结束标记，前端好判断本次对话已结束，允许下一轮对话
            chatWebSocketManager.sendOneMessage(userId, OVER_FLAG);
            return resultBuilder.toString();
        } catch (NoApiKeyException e) {
            log.error("调用通义千问缺少ApiKey");
            throw new ServiceException("调用通义千问缺少ApiKey");
        } catch (Exception e) {
            log.error("调用通义千问出现问题：{}", e.getMessage());
            throw new ServiceException("调用通义千问出现问题");
        }
    }

    private List<Message> buildMessageList(List<WxChatVo> wxChatVoList) {
        List<Message> messageList = new ArrayList<>();
        for (WxChatVo wxChatVo : wxChatVoList) {
            Message message = new Message();
            message.setRole(wxChatVo.getRole());
            message.setContent(wxChatVo.getContent());
            messageList.add(message);
        }
        return messageList;
    }

    public GenerationParam getBaseParam() {
        WxModelConfigVo config = getConfig();
        return GenerationParam.builder()
            .model(config.getModelType())
            .apiKey(config.getApiKey())
            .resultFormat(GenerationParam.ResultFormat.MESSAGE)
            .topP(0.8)
            .enableSearch(true)
            .incrementalOutput(true)
            .maxTokens(config.getMaxTokens())
            .build();
    }
}
